package hamster.csustef.warehouse.udf

import java.util

import org.apache.hadoop.hive.ql.exec.UDFArgumentException
import org.apache.hadoop.hive.ql.udf.generic.GenericUDTF
import org.apache.hadoop.hive.serde2.objectinspector.primitive.PrimitiveObjectInspectorFactory
import org.apache.hadoop.hive.serde2.objectinspector.{ObjectInspector, ObjectInspectorFactory, StructObjectInspector}

// 创建UDTF类继承GenericUDTF并重写下面的方法
// 生成自增数
// 传入两个数字，分别代表起始数字和自增量。如1234,23,则生成1234,1235,1236 ... 1256,一共23行数字。
class GenerateAutoIncrementUDTF extends GenericUDTF {

  override def initialize(argOIs: Array[ObjectInspector]): StructObjectInspector = {
    //判断传入的参数是否只有一个
    if (argOIs.length != 2) {
      throw new UDFArgumentException("有且只能有两个参数")
    }
    //判断参数类型
    if (argOIs(0).getCategory != ObjectInspector.Category.PRIMITIVE
      && argOIs(1).getCategory != ObjectInspector.Category.PRIMITIVE) {
      throw new UDFArgumentException("参数类型不匹配")
    }
    val fieldNames = new util.ArrayList[String]
    val fieldOIs = new util.ArrayList[ObjectInspector]
    fieldNames.add("auto_increment")
    fieldOIs.add(PrimitiveObjectInspectorFactory.javaStringObjectInspector)
    ObjectInspectorFactory.getStandardStructObjectInspector(fieldNames, fieldOIs)
  }

  override def process(objects: Array[AnyRef]): Unit = {
    var start: Int = objects(0).toString.toInt
    var plus: Int = objects(1).toString.toInt

    while (plus > 0) {
      val tmp = new Array[String](1)
      tmp(0) = start.toString
      forward(tmp)
      start += 1
      plus -= 1
    }
  }

  override def close(): Unit = {}
}
